#ifndef OPTION2D_HPP
#define OPTION2D_HPP

#include <iostream>
#include "Eigen/Dense"
#include "Eigen/Sparse"
using namespace std;
using namespace Eigen;

enum BoundCondiType
{
    Dirichlet,
    Neumann,
    Mixed
};

enum MultigridMethod
{
    Vcycle,
    FMG
};

namespace Operators
{
    enum Restriction
    {
        Injection,
        FullWeighting
    };

    enum Interpolation
    {
        linear,
        qudratic
    };
}

SparseMatrix<double> RestrictionMatrix(Operators::Restriction op, int n) // n = 2^i
{
    SparseMatrix<double> m(pow(n / 2 + 1, 2), pow(n + 1, 2));
    int cnt = 0;
    if (op == Operators::Restriction::Injection)
    {
        for (int i = 0; i <= n / 2; i++)
        {
            for (int j = 0; j <= n / 2; j++)
            {
                m.insert(cnt, 2 * i * (n + 1) + 2 * j) = 1;
                cnt++;
            }
        }
        m.finalize();
        // cout << m << endl;
    }
    else if (op == Operators::Restriction::FullWeighting)
    {
        for (int i = 0; i <= n / 2; i++)
        {
            m.insert(cnt, 2 * i) = 1;
            cnt++;
        }
        for (int i = 1; i <= n / 2 - 1; i++)
        {
            m.insert(cnt, 2 * i * (n + 1)) = 1;
            cnt++;
            for (int j = 1; j <= n / 2 - 1; j++)
            {
                m.insert(cnt, 2 * i * (n + 1) + 2 * j) = 0.5;
                m.insert(cnt, 2 * i * (n + 1) + 2 * j + 1) = 0.0625;
                m.insert(cnt, 2 * i * (n + 1) + 2 * j - 1) = 0.0625;
                m.insert(cnt, 2 * (i - 1) * (n + 1) + 2 * j - 1) = 0.0625;
                m.insert(cnt, 2 * (i - 1) * (n + 1) + 2 * j) = 0.0625;
                m.insert(cnt, 2 * (i - 1) * (n + 1) + 2 * j + 1) = 0.0625;
                m.insert(cnt, 2 * (i + 1) * (n + 1) + 2 * j) = 0.0625;
                m.insert(cnt, 2 * (i + 1) * (n + 1) + 2 * j + 1) = 0.0625;
                m.insert(cnt, 2 * (i + 1) * (n + 1) + 2 * j - 1) = 0.0625;
                cnt++;
            }
            m.insert(cnt, 2 * i * (n + 1) + n) = 1;
            cnt++;
        }
        for (int i = 0; i <= n / 2; i++)
        {
            m.insert(cnt, n * (n + 1) + 2 * i) = 1;
            cnt++;
        }
        m.finalize();
    }
    // cout << m << endl;
    return m;
}

SparseMatrix<double> InterpolationMatrix(Operators::Interpolation op, int n)
{
    SparseMatrix<double> M(pow(2 * n + 1, 2), pow(n + 1, 2));
    int cnt = 0;
    if (op == Operators::Interpolation::linear)
    {
        for (int i = 0; i <= n - 1; i++)
        {
            for (int j = 0; j <= n - 1; j++) //(i, j) (2*i, 2*j)
            {
                M.insert(2 * i * (2 * n + 1) + 2 * j, i * (n + 1) + j) = 1;

                M.insert(2 * i * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j) = 0.5;
                M.insert(2 * i * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 1) = 0.5;

                M.insert((2 * i + 1) * (2 * n + 1) + 2 * j, i * (n + 1) + j) = 0.5;
                M.insert((2 * i + 1) * (2 * n + 1) + 2 * j, (i + 1) * (n + 1) + j) = 0.5;

                M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j) = 0.25;
                M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j) = 0.25;
                M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 1) = 0.25;
                M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j + 1) = 0.25;
            }
            M.insert(2 * i * (2 * n + 1) + 2 * n, i * (n + 1) + n) = 1;
            M.insert((2 * i + 1) * (2 * n + 1) + 2 * n, i * (n + 1) + n) = 0.5;
            M.insert((2 * i + 1) * (2 * n + 1) + 2 * n, (i + 1) * (n + 1) + n) = 0.5;
            // cout << ++cnt << endl;
        }
        for (int j = 0; j <= n - 1; j++)
        {
            M.insert(2 * n * (2 * n + 1) + 2 * j, n * (n + 1) + j) = 1;
            M.insert(2 * n * (2 * n + 1) + 2 * j + 1, n * (n + 1) + j) = 0.5;
            M.insert(2 * n * (2 * n + 1) + 2 * j + 1, n * (n + 1) + j + 1) = 0.5;
        }
        M.insert(2 * n * (2 * n + 1) + 2 * n, n * (n + 1) + n) = 1;
        M.finalize();
    }
    else if (op == Operators::Interpolation::qudratic)
    {
        for (int i = 0; i <= n; i++)
        {
            for (int j = 0; j <= n; j++) //(i, j) (2*i, 2*j)
            {
                /*************格点************/
                M.insert(2 * i * (2 * n + 1) + 2 * j, i * (n + 1) + j) = 1;

                if (j <= n - 2)
                {
                    /*************左边横粗网格线上*************/
                    M.insert(2 * i * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j) = 0.375;
                    M.insert(2 * i * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 1) = 0.75;
                    M.insert(2 * i * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 2) = -0.125;
                }
                else if (j == n - 1)
                {
                    /*************最右侧横粗网格线上*************/
                    M.insert(2 * i * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j) = 0.75;
                    M.insert(2 * i * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 1) = 0.375;
                    M.insert(2 * i * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j - 1) = -0.125;
                }

                /*************纵粗网格线上*************/
                if (i <= n - 2)
                {
                    M.insert((2 * i + 1) * (2 * n + 1) + 2 * j, i * (n + 1) + j) = 0.375;
                    M.insert((2 * i + 1) * (2 * n + 1) + 2 * j, (i + 1) * (n + 1) + j) = 0.75;
                    M.insert((2 * i + 1) * (2 * n + 1) + 2 * j, (i + 2) * (n + 1) + j) = -0.125;
                }
                else if (i == n - 1)
                {
                    M.insert((2 * i + 1) * (2 * n + 1) + 2 * j, i * (n + 1) + j) = 0.75;
                    M.insert((2 * i + 1) * (2 * n + 1) + 2 * j, (i + 1) * (n + 1) + j) = 0.375;
                    M.insert((2 * i + 1) * (2 * n + 1) + 2 * j, (i - 1) * (n + 1) + j) = -0.125;
                }

                /*************左边粗网格之间*************/
                if (i <= n - 2)
                {
                    if (j <= n - 2)
                    {
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j) = 9.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 1) = 9.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 2) = -3.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j) = 9.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j + 1) = 9.0 / 16;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j + 2) = -3.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 2) * (n + 1) + j) = -3.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 2) * (n + 1) + j + 1) = -3.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 2) * (n + 1) + j + 2) = 1.0 / 64;
                    }
                    else if (j == n - 1)
                    {
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j) = 9.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 1) = 9.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j - 1) = -3.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j) = 9.0 / 16;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j + 1) = 9.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j - 1) = -3.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 2) * (n + 1) + j) = -3.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 2) * (n + 1) + j + 1) = -3.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 2) * (n + 1) + j - 1) = 1.0 / 64;
                    }
                }
                else if (i == n - 1)
                {
                    if (j <= n - 2)
                    {
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j) = 9.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 1) = 9.0 / 16;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 2) = -3.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i - 1) * (n + 1) + j) = -3.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i - 1) * (n + 1) + j + 1) = -3.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i - 1) * (n + 1) + j + 2) = 1.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j) = 9.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j + 1) = 9.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j + 2) = -3.0 / 64;
                    }
                    else if (j == n - 1)
                    {
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j) = 9.0 / 16;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j + 1) = 9.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, i * (n + 1) + j - 1) = -3.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i - 1) * (n + 1) + j) = -3.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i - 1) * (n + 1) + j + 1) = -3.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i - 1) * (n + 1) + j - 1) = 1.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j) = 9.0 / 32;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j + 1) = 9.0 / 64;
                        M.insert((2 * i + 1) * (2 * n + 1) + 2 * j + 1, (i + 1) * (n + 1) + j - 1) = -3.0 / 64;
                    }
                }

                // cout << ++cnt << endl;
            }
            M.finalize();
        }
    }
    //cout << M << endl;
    return M;
}

#endif